import os
import random
from PIL import Image
import numpy as np
import torch
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score, roc_auc_score
from flair import FLAIRModel
from tqdm import tqdm


# =====================
# 1. 数据加载函数
# =====================

def load_balanced_dataset(root_dir, main_class="1", total_target=5000):
    """
    从root_dir中加载数据：
    - 主类(main_class)全部加载
    - 其他类随机采样，数量均衡，总数约为total_target
    """
    images, labels = [], []

    # 找到所有类别文件夹
    all_classes = [d for d in os.listdir(root_dir) if os.path.isdir(os.path.join(root_dir, d))]
    all_classes.sort()

    # 主类路径
    main_dir = os.path.join(root_dir, main_class)
    main_imgs = [os.path.join(main_dir, f) for f in os.listdir(main_dir)
                 if f.lower().endswith((".png", ".jpg", ".jpeg"))]
    n_main = len(main_imgs)
    print(f"[INFO] 主类 {main_class}: {n_main} 张")

    # 剩余数量 & 其他类数量
    remaining = max(total_target - n_main, 0)
    other_classes = [c for c in all_classes if c != main_class]
    n_other = len(other_classes)
    per_class = remaining // n_other if n_other > 0 else 0
    print(f"[INFO] 每个非主类采样约 {per_class} 张")

    # 主类全部加载
    for path in main_imgs:
        try:
            img = Image.open(path).convert("RGB")
            images.append(np.array(img))
            labels.append(int(main_class))
        except Exception as e:
            print(f"[WARN] 跳过无法加载的图像 {path}: {e}")

    # 其他类随机采样
    for c in other_classes:
        c_dir = os.path.join(root_dir, c)
        all_imgs = [os.path.join(c_dir, f) for f in os.listdir(c_dir)
                    if f.lower().endswith((".png", ".jpg", ".jpeg"))]
        if len(all_imgs) == 0:
            continue
        sampled = random.sample(all_imgs, min(per_class, len(all_imgs)))
        print(f"[INFO] 类别 {c}: 采样 {len(sampled)} 张")

        for path in sampled:
            try:
                img = Image.open(path).convert("RGB")
                images.append(np.array(img))
                labels.append(int(c))
            except Exception as e:
                print(f"[WARN] 跳过无法加载的图像 {path}: {e}")

    # 打乱顺序
    combined = list(zip(images, labels))
    random.shuffle(combined)
    images, labels = zip(*combined)

    print(f"[INFO] 最终加载总数: {len(images)} 张")
    return list(images), np.array(labels)

def load_binary_balanced_dataset(root_dir, total_target=5000):
    """
    从 root_dir 中加载二分类数据 (0 和 1 文件夹)：
    - 尝试从两个类别中加载相同比例的数据
    - 若某类图片不足，则按实际数量加载
    """
    images, labels = [], []

    class_dirs = ["0", "1"]
    class_counts = {}

    # 统计每类图像数量
    for c in class_dirs:
        c_dir = os.path.join(root_dir, c)
        if not os.path.exists(c_dir):
            print(f"[WARN] 类别文件夹 {c_dir} 不存在，跳过。")
            continue
        imgs = [os.path.join(c_dir, f) for f in os.listdir(c_dir)
                if f.lower().endswith((".png", ".jpg", ".jpeg"))]
        class_counts[c] = imgs
        print(f"[INFO] 类别 {c}: 共 {len(imgs)} 张")

    # 计算每类应加载数量
    per_class = total_target // 2
    print(f"[INFO] 每类目标加载 {per_class} 张")

    for c, img_list in class_counts.items():
        n_available = len(img_list)
        n_sample = min(per_class, n_available)
        sampled = random.sample(img_list, n_sample)
        print(f"[INFO] 类别 {c}: 实际加载 {len(sampled)} 张")

        for path in sampled:
            try:
                img = Image.open(path).convert("RGB")
                images.append(np.array(img))
                labels.append(int(c))
            except Exception as e:
                print(f"[WARN] 跳过无法加载的图像 {path}: {e}")

    # 打乱顺序
    combined = list(zip(images, labels))
    random.shuffle(combined)
    images, labels = zip(*combined)

    print(f"[INFO] 最终加载总数: {len(images)} 张")
    print(f"[INFO] 类别分布: 0={np.sum(np.array(labels)==0)}, 1={np.sum(np.array(labels)==1)}")
    return list(images), np.array(labels)

def infer_one():
    # =====================
    # 2. 加载数据
    # =====================
    data_folder = "/media/disk/02drive/13hias/02yixuan/ZSL/data/eyepacs/train" 
    images, labels = load_binary_balanced_dataset(data_folder, total_target=5000)

    # =====================
    # 3. 加载模型
    # =====================
    model = FLAIRModel(
        from_checkpoint=True,
        bert_type="/media/disk/02drive/13hias/02yixuan/ZSL/Bio_ClinicalBERT",
        weights_path="/media/disk/02drive/13hias/02yixuan/ZSL/flair/modeling/flair_pretrained_weights/flair_resnet.pth"
    )

    text_prompts = ["normal", "healthy", "macular edema", "diabetic retinopathy", "glaucoma", "macular hole",
        "lesion", "lesion in the macula"]

    preds = []
    dr_labels = {"diabetic retinopathy"}

    for img_np in tqdm(images, desc="Inferencing"):
        probs, logits = model(img_np, text_prompts)
        pred_idx = int(np.argmax(probs))
        pred_label = text_prompts[pred_idx]
        pred = 1 if pred_label in dr_labels else 0
        preds.append(pred)

    preds = np.array(preds)

    # =====================
    # 5. 计算评估指标
    # =====================
    acc = accuracy_score(labels, preds)
    prec = precision_score(labels, preds, average="macro")
    rec = recall_score(labels, preds, average="macro")
    f1 = f1_score(labels, preds, average="macro")

    try:
        if len(np.unique(labels)) > 2:
            auc = roc_auc_score(labels, torch.nn.functional.one_hot(torch.tensor(preds), num_classes=len(np.unique(labels))), multi_class="ovo")
        else:
            auc = roc_auc_score(labels, preds)
    except Exception:
        auc = np.nan

    # ✅ 返回指标字典
    return {
        "Accuracy": acc,
        "Precision": prec,
        "Recall": rec,
        "F1": f1,
        "AUC": auc
    }


# =====================
# 主循环：多次推理并汇总结果
# =====================
all_metrics = []

for i in range(10):
    print(f"\n===== 第 {i+1} 次推理 =====")
    result = infer_one()
    all_metrics.append(result)

print("\n========== 所有实验结果 ==========")
for i, m in enumerate(all_metrics, 1):
    print(f"Run {i}: {m}")

# ✅ 如果需要统计平均值
mean_metrics = {k: np.mean([m[k] for m in all_metrics]) for k in all_metrics[0]}
print("\n========== 平均指标 ==========")
for k, v in mean_metrics.items():
    print(f"{k}: {v:.4f}")
